package com.wangwg.sparkTest.PostGIS

import java.io.File
import java.util.UUID

import org.apache.hadoop.conf.Configuration
import org.apache.spark.{SparkConf, SparkContext, TaskContext}
import org.geotools.data.Query
import org.locationtech.geomesa.spark.{GeoMesaSpark, GeoMesaSparkKryoRegistrator, SpatialRDD}

import scala.collection.JavaConversions._


/**
 * 打开postgis数据库，遍历循环输出数据
 * 数据读取测试例子
 *
 * @author wangwg
 */
object GeomesaPostGISRead {
  def main(args: Array[String]): Unit = {
    val sparkConf = new SparkConf().setMaster("local[*]").setAppName("GeomesaPostGISRead");
    sparkConf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
    sparkConf.set("spark.kryo.registrator", classOf[GeoMesaSparkKryoRegistrator].getName)
    val sparkContext = SparkContext.getOrCreate(sparkConf);
    val params = Map(
      "geotools" -> "true",
      "dbtype" -> "postgis",
      "host" -> "139.9.130.25",
      "user" -> "lhpostgis",
      "passwd" -> "lhpostgis",
      "port" -> "5432",
      "database" -> "lhpostgis");
    val query = new Query("SJ_KCZY")
    val rdd = GeoMesaSpark(params).rdd(new Configuration(), sparkContext, params, query)
    println(rdd.schema) //打印数据存储方案
    printFeature(sparkContext, rdd);
    println("end");
  }

  def printFeature(sparkContext: SparkContext, rdd: SpatialRDD): Unit = {
    println("默认分区数:".concat(rdd.getNumPartitions.toString))
    //val reRdd = rdd.repartition(3) //将rdd重新分为5个分区
    val reRdd = rdd //将rdd重新分为5个分区
    println("现有分区数:".concat(reRdd.getNumPartitions.toString))
    reRdd.foreachPartition { iter =>
      import java.nio.charset.Charset

      import org.apache.commons.io.FileUtils

      //println("当前分区id:".concat(TaskContext.getPartitionId.toString).concat(",数据量:").concat(iter.length.toString)) 在这打印后，foreach next是没有数据的，下面循环无法执行
      iter.foreach(item => {
        var path = "D:\\work\\bigdata\\".concat(UUID.randomUUID().toString).concat(".txt");
        val file = new File(path)
        val info = "当前分区id:".concat(TaskContext.getPartitionId.toString).concat(":").concat(item.toString)
        FileUtils.write(file, info, Charset.forName("UTF-8"))
        //println("当前分区id:".concat(TaskContext.getPartitionId.toString).concat(":").concat(item.toString))
      }
      )
    }
  }
}